import pylab as P
import matplotlib as M
import numpy as N
from scipy.signal import detrend
from scipy.stats import scoreatpercentile

lineprops = dict(linewidth=1, color='black', linestyle='-')
lineprops2 = dict(linewidth=1, color='red', linestyle='-')

global fig_main,ax_main,lines_main,scale_bar_main, \
    scale_text_main,signals_main,overlap_main, t_main

def _get_max(signals):
    max_mean = scoreatpercentile(N.array([N.abs(sig).max() for sig in signals]), \
        per = 75)
    return max_mean

def _PlotInitSignals(t,signals,ylabels,title,overlap):
    global lineprops
    P.ioff()
    fig = P.figure()
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])
    ax.set_xlabel('time in ms')
    if title: ax.set_title(title)
    #get 75 percentile maximum values of channels to determine overlap
    max_val = _get_max(signals)
    mean_locs = 1 + N.arange(len(signals))*2 * (1-overlap)
    tick_locs = 1 + N.arange(len(signals))*2 * (1-overlap) \
    + N.array(signals)[:,0] / max_val
    lines = []
    for i, sig in enumerate(signals):
        if i%2==0:
            lines.append(ax.plot(t,sig/max_val + 1 + 2*i * (1-overlap), **lineprops))
        else: lines.append(ax.plot(t,sig/max_val + 1 + 2*i * (1-overlap), **lineprops2))
    ax.set_yticks(tick_locs)
    ax.set_yticklabels(ylabels)
    ax.set_ylim(mean_locs[0]-1.5+overlap, mean_locs[-1]+1.5-overlap)
    ax.set_xlim(t[0],t[-1])
    ylims = ax.get_ylim()
    yrange = (N.abs(ylims)).sum()
    scale_bar = ax.axvline(t[int(len(t)*0.9)], ymin = 0.03, \
    ymax = 0.03 + 1/float(yrange), **lineprops)
    scale_text = ax.text(t[int(len(t)*0.92)],ylims[0]+ (0.03 + 0.5/float(yrange))*yrange, \
    '%.2f' % max_val, verticalalignment='center')
    return fig, ax, lines, scale_bar, scale_text

def _UpdatePlot(fig,ax,lines,scale_bar,scale_text,t,signals,overlap):
    global lineprops
    P.ioff()
    max_val = _get_max(signals)
    mean_locs = 1 + N.arange(len(signals))*2 * (1-overlap)
    tick_locs = 1 + N.arange(len(signals))*2 * (1-overlap) \
    + N.array(signals)[:,0] / max_val
    for i, line in enumerate(lines):
        line[0].set_ydata(signals[i]/max_val + 1 + 2*i * (1-overlap))
    ax.set_yticks(tick_locs)
    ax.set_ylim(mean_locs[0]-1.5+overlap, mean_locs[-1]+1.5-overlap)
    ax.set_xlim(t[0],t[-1])
    ylims = ax.get_ylim()
    yrange = (N.abs(ylims)).sum()
    scale_bar.set_ydata(N.array([0.03, 0.03 + 1/float(yrange)]))
    scale_text.set_position(N.array([t[int(len(t)*0.92)],ylims[0]+ (0.03 + \
    0.5/float(yrange))*yrange]))
    scale_text.set_text('%.2f' % max_val)
    P.draw()
    return

def PlotEEG(signals, ylabels, t, overlap=0.5, title = False):
    """Function to Plot EEG-Signals in of several channels
    Def: PlotEEG(signals, ylabels, t, overlap=0.5, title = False)

    Scale can be adjusted by pressing \'+\' or \'-\'

    --input:
        -signals: list of input data
              signal[0]: first complete channel
              signal[1]: second complete channel
              .
              .
        -ylabels: list of channel names in same order as in \'signals\'
        -t: array of time values
        -overlap: scale value in interval [0,1]
              for 0, signals don't overlap at all
              for 1, signal's means fall onto each other
        -title: title string
    --output:
        -fig: figure instance, plot with fig.plot()
    """
    global fig_main,ax_main,lines_main,scale_bar_main, \
    scale_text_main,signals_main,overlap_main, t_main
    n_plots= len(signals)
    signals = [detrend(sig, type='linear') for sig in signals]
    fig, ax, lines, scale_bar, scale_text = _PlotInitSignals(t,signals,ylabels,overlap=overlap, title = title)
    fig_main = fig
    ax_main = ax
    lines_main = lines
    scale_bar_main = scale_bar
    scale_text_main = scale_text
    overlap_main = overlap
    t_main = t
    signals_main = signals
    fig.text(0.5, 0.02, 'Click on axis and press \'+\' or \'-\' to adjust scale', ha='center', \
    va='baseline')
    P.connect('key_press_event', _keypress)
    return fig

def _adjust_overlap(overlap,direction):
    x = N.tan(overlap * N.pi - N.pi/2.)
    x += 0.25*direction
    overlap = (N.arctan(x) + N.pi/2) / N.pi
    return overlap

def _keypress(event):
    global fig_main,ax_main,lines_main,scale_bar_main, \
    scale_text_main,signals_main,overlap_main, t_main
    if event.key in ('+', '='):
            overlap_main = _adjust_overlap(overlap_main,1)
    elif event.key in ('-', '_'):
            overlap_main = _adjust_overlap(overlap_main,-1)
    _UpdatePlot(fig_main,ax_main,lines_main,scale_bar_main, \
    scale_text_main,t_main,signals_main,overlap_main)
    return
